/**
	@file
	Paging and basic memory management (frames level) routines.
	@note Written for JamesM's kernel development tutorials.
	@todo Complete and register interrupt handler.
	@par Edited:
		Konstantin Kastanov
*/
// paging.c -- Defines the interface for and structures relating to paging.
//             Written for JamesM's kernel development tutorials.

#include "headers/paging.h"
#include "headers/kheap.h"
#include "../adt/headers/heap.h"
#include "../tools/headers/common.h"
#include "../api/headers/string.h"
#include "../api/headers/io.h"
#include "../interrupts/headers/idt.h"

#define LOADING_DOTS_N 9

void page_fault_handler(registers_type registers);
static page_table_t *clone_table(page_table_t *src, u32int *physAddr);
extern volatile BOOL critical_section;

/** The kernel's page directory*/
page_directory_t *kernel_directory=0;

/** The current page directory;*/
page_directory_t *current_directory=0;

/** A bitset of frames - used or free.*/
u32int *frames;
/** Number of frames*/
u32int nframes;

// Defined in kheap.c
extern u32int placement_address;
extern heap_t *kheap;
// Macros used in the bitset algorithms.
/** Gets bitset index (\ref frames array)*/
#define INDEX_FROM_BIT(a) (a/(8*4))
/** Gets offset in a specified (returned by \ref INDEX_FROM_BIT macro) u32int in \ref frames array*/
#define OFFSET_FROM_BIT(a) (a%(8*4))

// Sets a bit in the frames bitset
static void set_frame(u32int frame_addr)
{
	
	BOOL backup =critical_section;
	critical_section = true;
    u32int frame = frame_addr;
    u32int idx = INDEX_FROM_BIT(frame);
    u32int off = OFFSET_FROM_BIT(frame);
    frames[idx] |= (0x1 << off);
	critical_section = backup;
}

// Clears a bit in the frames bitset
static void clear_frame(u32int frame_addr)
{
	BOOL backup = critical_section;
	critical_section = true;
    u32int frame = frame_addr;
    u32int idx = INDEX_FROM_BIT(frame);
    u32int off = OFFSET_FROM_BIT(frame);
    frames[idx] &= ~(0x1 << off);
	critical_section = backup;
}

// Tests if a bit is set.
static u32int test_frame(u32int frame_addr)
{
    u32int frame = frame_addr;
    u32int idx = INDEX_FROM_BIT(frame);
    u32int off = OFFSET_FROM_BIT(frame);
    return (frames[idx] & (0x1 << off));
}

// Finds the first free frame. Optimized algorithm
static u32int first_frame()
{
    u32int i, j;
    for (i = 0; i < INDEX_FROM_BIT(nframes); i++)
    {
        if (frames[i] != 0xFFFFFFFF) // nothing free, exit early.
        {
            // at least one bit is free here.
            for (j = 0; j < 32; j++)
            {
                u32int toTest = 0x1 << j;
                if ( !(frames[i]&toTest) )
                {
                    return i*4*8 + j;
                }
            }
        }
    }
}

/** Allocates a first free frame.
	@param page Page for which to find place
	@param is_kernel True if page is in kernel mode
	@param is_writable True if page is writable (only applicable to user mode, have no effect in kernel mode, unless a flag in CR0 is set)
	@todo Add panic handler
	@return 0 if successfull (or page is already allocated), MEMORY_ERROR- otherwise. 
*/
u32int alloc_frame(page_t *page, int is_kernel, int is_writeable)
{
	setLastError(NO_ERROR);
    if (page->frame != 0) //Already allocated?
    {
        return 0;
    }
    if(!page)
	{
		setLastError(ERR_INVALID_ARG);
		return MEMORY_ERROR;
	}
    else
    {
        u32int idx = first_frame();
        if (idx == (u32int)-1)
        {
            setLastError(ERR_OUT_OF_MEMORY);
			return MEMORY_ERROR;
        }
        set_frame(idx);
        page->present = 1;
        page->rw = (is_writeable)?1:0;
        page->user = (is_kernel)?0:1;
        page->frame = idx;
    }
    return 0;
}

/**
	Deallocates a speciified frame.
	@param page Page, that occupies the frame to be freed
 */
void free_frame(page_t *page)
{
    u32int frame;
	setLastError(NO_ERROR);
    if (!(frame=page->frame)) //Frame was not allocated
    {
        return;
    }
    else
    {
        clear_frame(frame);
        page->frame = 0x0;
    }
}

/**
	Sets up environment and enables paging.
	@todo Pass the size of physical memory from main functions
 */
void initialise_paging(u32int memory_size)
{
#ifdef DEBUG
	printf("PA1: %h\n", placement_address);
#endif
	setLastError(NO_ERROR);
    // The size of physical memory. For the moment we 
    // assume it is 16MB big.
    u32int mem_end_page = memory_size;//0x1000000; //Should be passed as parameter!
    
    nframes = mem_end_page / 0x1000;
    frames = (u32int*)kmalloc(INDEX_FROM_BIT(nframes));
	if(getLastError() != NO_ERROR)
	{
		K_PANIC("Error allocating memory for frames bitmap");
	}
    memset(frames, 0, INDEX_FROM_BIT(nframes));
    
    // Let's make a page directory.
	u32int phys;
    kernel_directory = (page_directory_t*)kmalloc_a(sizeof(page_directory_t));
	memset(kernel_directory, 0, sizeof(page_directory_t));
	kernel_directory->physicalAddr = (u32int)kernel_directory->tablesPhysical;
    //current_directory = kernel_directory;
#ifdef DEBUG
	printf("PA2: %h\n", placement_address);
#endif
	// Map some pages in the kernel heap area.
    // Here we call get_page but not alloc_frame. This causes page_table_t's 
    // to be created where necessary. We can't allocate frames yet because they
    // they need to be identity mapped first below, and yet we can't increase
    // placement_address between identity mapping and enabling the heap!
    int i = 0;
    for (i = KHEAP_START; i < KHEAP_START + KHEAP_INITIAL_SIZE; i += 0x1000) //Allocate memory for heap page tables
	{
		get_page(i, 1, kernel_directory);
		if(getLastError() != NO_ERROR)
		{
			K_PANIC("Error allocating memory for heap page tables");
		}
	}
	
	
	
#ifdef DEBUG
	printf("PA3: %h\n", placement_address);
#endif
    // We need to identity map (phys addr = virt addr) from
    // 0x0 to the end of used memory, so we can access this
    // transparently, as if paging wasn't enabled.
    // NOTE that we use a while loop here deliberately.
    // inside the loop body we actually change placement_address
    // by calling kmalloc(). A while loop causes this to be
    // computed on-the-fly rather than once at the start.
    i = 0;
	printf("Allocating pages.");
	u32int upper_address = (u32int)(placement_address + 0x1000 + ((u32int)(placement_address / 0x1000) * sizeof(page_table_t)));
	if(upper_address & 0xFFF)
		upper_address = upper_address & 0xFFFFF000 + 0x1000;
	u32int dots_rate = (u32int)upper_address / LOADING_DOTS_N;
	u32int count = 0;
		
    while (i < placement_address + 0x1000) //Allocate frames for kernel code and initially allocated memory (plus one page for heap allocations)
    {
        // Kernel code is readable but not writeable from userspace.
        
		if(alloc_frame(get_page(i, 1, kernel_directory), 0, 0) == MEMORY_ERROR)
		{
			K_PANIC("Error allocating frames for kernel code");
		}
		if(i >= count * dots_rate)
		{
			printf(".");
			count++;
		}
		//printf("(%h : %h)  ", i, i);
        i += 0x1000;
    }
    printf("\n");
#ifdef DEBUG
    printf("PA3: %h\n", placement_address);
#endif
    
     // Now allocate those pages we mapped earlier.
    for (i = KHEAP_START; i < KHEAP_START + KHEAP_INITIAL_SIZE; i += 0x1000)
	{
		if(alloc_frame( get_page(i, 1, kernel_directory), 0, 0) == MEMORY_ERROR)
		{
			K_PANIC("Error allocating frames for heap");
		}
	}
	
#ifdef DEBUG
	printf("PA2: %h\n", placement_address);
#endif
    // Before we enable paging, we must register our page fault handler.
   register_interrupt_handler(14, page_fault_handler);

    // Now, enable paging!
    switch_page_directory(kernel_directory);
	
	kheap = create_heap(KHEAP_START, KHEAP_START + KHEAP_INITIAL_SIZE, KHEAP_START + KHEAP_MAX_SIZE, kernel_directory, 0, 0);
	if(getLastError() != NO_ERROR)
	{
		K_PANIC("Error creating kernel heap");
	}
	
	current_directory = clone_directory(kernel_directory);
	
	if(getLastError() != NO_ERROR)
	{
		K_PANIC("Error cloning kernel directory");
	}
	
	//current_directory = kernel_directory;
#ifdef DEBUG
	printf("kernel: 0x%h, 0x%h\nclone: 0x%h, 0x%h\n", kernel_directory->tablesPhysical[0],
			kernel_directory->tablesPhysical[768], current_directory->tablesPhysical[0],
			current_directory->tablesPhysical[768]);


	printf("Old physical: 0x%h\n", kernel_directory->physicalAddr);
	printf("New physical: 0x%h\n", current_directory->physicalAddr);
#endif
	//while(1);
	switch_page_directory(current_directory);
	
}

/**
	Loads specified page's address into cr3 register and enables paging.
	@param dir New page directory
 */
void switch_page_directory(page_directory_t *dir)
{
	asm volatile("cli");
    current_directory = dir;
    asm volatile("mov %0, %%cr3":: "r"(dir->physicalAddr));
    u32int cr0;
    asm volatile("mov %%cr0, %0": "=r"(cr0));
    cr0 |= 0x80000000; // Enable paging!
   
    asm volatile("mov %0, %%cr0":: "r"(cr0));
	asm volatile("sti");
}

/**
	Gets page struct from a specified page directory for a specified address.
	@param address Specified address
	@param make If the required page table is not found (NULL) and make == true, than page table will be created.
 */
page_t *get_page(u32int address, BOOL make, page_directory_t *dir)
{
    // Turn the address into an index.
    address /= 0x1000;
    // Find the page table containing this address.
    u32int table_idx = address / 1024;
	setLastError(NO_ERROR);
    if (dir->tables[table_idx]) // If this table is already assigned
    {
        return &dir->tables[table_idx]->pages[address%1024];
    }
    else if(make)
    {
        u32int tmp;
        dir->tables[table_idx] = (page_table_t*)kmalloc_ap(sizeof(page_table_t), &tmp);
		if(getLastError() != NO_ERROR)
			return 0;
		memset(dir->tables[table_idx], 0, 0x1000);
        dir->tablesPhysical[table_idx] = tmp | 0x7; // PRESENT, RW, US.
        return &dir->tables[table_idx]->pages[address%1024];
    }
    else
    {
        return 0;
    }
}


void page_fault_handler(registers_type registers)
{
    // A page fault has occurred.
    // The faulting address is stored in the CR2 register.
    u32int faulting_address;
    asm volatile("mov %%cr2, %0" : "=r" (faulting_address));
    
    // The error code gives us details of what happened.
    int present   = !(registers.error_code & 0x1); // Page not present
    int rw = registers.error_code & 0x2;           // Write operation?
    int us = registers.error_code & 0x4;           // Processor was in user-mode?
    int reserved = registers.error_code & 0x8;     // Overwritten CPU-reserved bits of page entry?
    int id = registers.error_code & 0x10;          // Caused by an instruction fetch?

    // Output an error message.
    printf("Page fault! ( ");
    if (present) {printf("present ");}
    if (rw) {printf("read-only ");}
    if (us) {printf("user-mode ");}
    if (reserved) {printf("reserved ");}
    printf(") at 0x");
    printf("%h ", faulting_address);
	printf("Error code: %d", registers.error_code);
    printf("\n");
    //K_PANIC("Page fault");
}

//Returns 0 on error
static page_table_t *clone_table(page_table_t *src, u32int *physAddr)
{
	setLastError(NO_ERROR);
    // Make a new page table, which is page aligned.
    page_table_t *table = (page_table_t*)kmalloc_ap(sizeof(page_table_t), physAddr);
	if(getLastError() != NO_ERROR)
		return 0;
    // Ensure that the new table is blank.
    memset(table, 0, sizeof(page_table_t));

    // For every entry in the table...
    int i;
    for (i = 0; i < 1024; i++)
    {
        // If the source entry has a frame associated with it...
        if (src->pages[i].frame)
        {
            // Get a new frame.
            if(alloc_frame(&table->pages[i], 0, 0) == MEMORY_ERROR)
				return 0;
			
            // Clone the flags from source to destination.
            if (src->pages[i].present) table->pages[i].present = 1;
            if (src->pages[i].rw) table->pages[i].rw = 1;
            if (src->pages[i].user) table->pages[i].user = 1;
            if (src->pages[i].accessed) table->pages[i].accessed = 1;
            if (src->pages[i].dirty) table->pages[i].dirty = 1;
            // Physically copy the data across. This function is in process.s.
            copy_page_physical(src->pages[i].frame * 0x1000, table->pages[i].frame * 0x1000);
        }
    }
    return table;
}

//Returns 0 at error
page_directory_t *clone_directory(page_directory_t *src)
{
    u32int phys;
	setLastError(NO_ERROR);
    // Make a new page directory and obtain its physical address.
    page_directory_t *dir = (page_directory_t*)kmalloc_ap(sizeof(page_directory_t), &phys);
	if(getLastError() != NO_ERROR)
		return 0;
    // Ensure that it is blank.
    memset(dir, 0, sizeof(page_directory_t));

    // Get the offset of tablesPhysical from the start of the page_directory_t structure.
    u32int offset = (u32int)dir->tablesPhysical - (u32int)dir;

    // Then the physical address of dir->tablesPhysical is:
    dir->physicalAddr = phys + offset;
    // Go through each page table. If the page table is in the kernel directory, do not make a new copy.
    int i,a;
    for (i = 0; i < 1024; i++)
    {
        if (!src->tables[i])
            continue;

        if (kernel_directory->tables[i] == src->tables[i])
        {
            // It's in the kernel, so just use the same pointer.
            dir->tables[i] = src->tables[i];
            dir->tablesPhysical[i] = src->tablesPhysical[i];
        }
        else
        {
            // Copy the table.
            u32int phys;
            dir->tables[i] = clone_table(src->tables[i], &phys);
			if(getLastError() != NO_ERROR)
				return 0;
            dir->tablesPhysical[i] = phys | 0x07;
        }
    }
    return dir;
}

